from fastapi.security import OAuth2PasswordBearer
from jose import JWTError, jwt
from datetime import datetime
from database import get_db
from fastapi import Depends, HTTPException, status
from config import SECRET_KEY, ALGORITHM
from fastapi import Request
import asyncio
from .llm import OllamaCli
from models import User


oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")


async def get_current_user(token: str = Depends(oauth2_scheme), db=Depends(get_db)):
    # 验证登录
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
    except JWTError:
        raise credentials_exception
    
    user = db.query(User).filter(User.username == username).first()
    if user is None:
        raise credentials_exception
    return user


async def event_generator(request: Request):
    while True:
        if await request.is_disconnected():
            break
        now = datetime.datetime.now()
        yield f"data: {now.timestamp()}\n\n"
        await asyncio.sleep(1)

async def event_generator_ollama(message: str):
    cli = OllamaCli()
    messages = [
        {"role": "system", "content": "你现在是一名合格的植保员，请扮演好您的角色。"},
        {"role": "user", "content": message}
    ]
    
    response = await cli.async_chat(messages)
    async for res in response:       
        yield f"{res}\n"